import numpy as np
import PIL.Image


def load_spe(name: str, samples: int, lines: int, bands: int, data_type: int):
    """
    加载 SPE 文件, 只支持 BIL 格式, 其他格式参考文档自行实现
    文件格式查看 hdr 文件中的 interleave 字段

    :param name: SPE 文件名
    :param samples: 空间通道数, 对应 hdr 文件中的 samples
    :param lines: 帧数, 对应 hdr 文件中的 lines
    :param bands: 光谱通道数, 对应 hdr 文件中的 bands
    :param data_type: 数据类型, 对应 hdr 文件中的 data type, 表示每个信号值的数据类型
    :return: np.memmap
    """
    if data_type == 0:
        dtype = np.uint8
    elif data_type == 4:
        dtype = np.float32
    elif data_type == 12:
        dtype = np.uint16
    else:
        raise AssertionError('不支持的数据类型')
    return np.memmap(name, dtype, 'r', shape=(lines, bands, samples))


def preview(raw: np.ndarray, bands: list[int], thresholds: list[float]):
    """
    生成预览图
    :param raw: 原始信号数据
    :param bands: 空间通道, len(bands) == 1 or len(bands) == 3, len(bands) == 1 时返回灰度图, len(bands) == 3 时返回 RGB 图
    :param thresholds: 阈值, len(thresholds) == len(bands)
    :return: np.ndarray, 只有一个空间通道时 shape => (height, width), 三个空间通道时 shape => (height, width, 3)
    """
    assert len(bands) == 1 or len(bands) == 3, f'通道数必须为 1 或 3, 当前: {len(bands)}'
    assert len(bands) == len(thresholds), '阈值错误'
    im = np.clip(np.array(raw[:, bands, :]).transpose(2, 0, 1) / thresholds * 255, 0, 255).astype(np.uint8)
    if len(bands) == 1:
        im = im.reshape((raw.shape[2], -1),)
    return im


if __name__ == '__main__':
    data = load_spe(r'PET-HDPE.spe', 320, 319, 256, 12)
    # 灰度图
    # gray = PIL.Image.fromarray(preview(data, [226], [800]), 'L')
    # gray.show()

    # RGB 图
    rgb = PIL.Image.fromarray(preview(data, [226, 124, 84], [800, 1100, 400]), 'RGB')
    rgb.show()
